#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/random.h>
#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include <iostream>
#include <cstdlib>
#include <ctime>
#include <fstream>
#include <iomanip>
#define N 2048
#define L0 10
#define V0 1
#define speed 0.2
#define Steps 10
#define  K0 4 
using namespace std;

void print(const thrust::device_vector< thrust::pair<int,int> >& v)
{
	for(size_t i = 0; i < v.size(); i++)
	{
		thrust::pair<int,int> p = v[i];
		std::cout << " (" << p.first << "," << p.second << ")";
	}
	std::cout << "\n";
}
void print(const thrust::device_vector<int>&v)
{
	for(size_t i = 0; i < v.size(); i++)
	{
		cout<<setw(2)<<v[i]<<" ";
	}
	std::cout << "\n";
}
void print(const thrust::device_vector<double>&v)
{
	for(size_t i = 0; i < v.size(); i++)
	{
		cout<<v[i]<<" ";
	}
	std::cout << "\n";
}
__device__  void cuda_findKthMin2(const double* data, int n,int K,
	int* id)
{
	for (int i = 0; i < n; i++)
		id[i] = i;
	if ( K >=n ) return ;
	bool flag = false;
	int k = K;
	int low = 0;
	int high = n - 1;
	while ( !flag )
	{
		int key = id[low];
		int x = low;
		int y = high;
		while ( x < y )
		{
			while( data[key] <= data[id[y]] && x < y) y--;
			id[x] = id[y];
			while( data[id[x]] <= data[key] && x < y) x++;
			id[y] = id[x];
		}
		id[x] = key;
		int cnt = x - low + 1;
		if (k == cnt ) flag = true;
		if ( k > cnt ) { k = k - cnt; low = x + 1;} 
		else { high = x - 1; }
	}
}

__global__ void cuda_hash_mapping(double* dx, int* d_cell, int n,
								  double a, double b, int CELLNUM)
{
	int tid = threadIdx.x + blockIdx.x * blockDim.x;
	if (tid < n)
	{
		double x = dx[tid];
		double cellSize = (b - a) / CELLNUM;
		int ID = ( x - a ) / cellSize; //error checking might be added
		d_cell[tid] = ID;
	}
}

void cuda_sort_cell(thrust::device_vector<int> &cell, thrust::device_vector<int> &ID)
{
	for(size_t i = 0; i < cell.size(); i++)
		ID[i] = i;
	thrust::device_vector< thrust::pair<int,int> > pairs( cell.size() );
	for (size_t i = 0; i < pairs.size(); i++)
		pairs[i] = thrust::make_pair(cell[i], ID[i]);
	thrust::sort(pairs.begin(), pairs.end());
	for (size_t i = 0; i < pairs.size(); i++)
	{
		thrust::pair<int,int> p = pairs[i];
		cell[i] = p.first;
		ID[i] = p.second;
		//cout<<" ("<<cell[i]<<","<<ID[i]<<")"<<endl;
	}
}

__global__ void cuda_count_cell(int* d_cell, int* d_start, int n, int cellnum)
{
	int tid = threadIdx.x + blockIdx.x * blockDim.x;
	if( tid < n )
	{
		if ( tid == 0 ) {
			d_start[tid] = 0;
		}else {
			if ( d_cell[tid] != d_cell[tid - 1])
				d_start[d_cell[tid]] = tid;
		}
	}
	if ( tid == n)
		d_start[cellnum] = n;
}

__device__ int cuda_get_cell_size(int* cell, int* cell_start, 
								   int cell_ID, int n, int cell_num)
{
	int next_cell = cell_ID + 1;
	while( cell_start[next_cell] == 0 && next_cell != cell_num  ) next_cell++; 

	return (cell_start[next_cell] - cell_start[cell_ID] );
}

__device__ void cuda_get_cell_size(int* cell, int* cell_start, 
								   double x, int n, int cell_num,
								   int& cell_size)
{
	double cell_length = ( L0 - 0 ) / cell_num;
	int cell_ID = (x - 0 ) / cell_length;
	int next_cell = cell_ID + 1;

	while( cell_start[next_cell] == 0 && next_cell != cell_num  ) next_cell++; 

	cell_size = cell_start[next_cell] - cell_start[cell_ID];
}

__device__ void cuda_get_cell_list( int* cell, int* cell_start, double x,
									int n, int k, int cell_num, 
									int** cell_list, int& cell_list_size, int& count)
{
	//by default, cell_list is a list of size cell_num, though some of the entries are not used.
	int size;
	double cell_length = ( L0 - 0 ) / cell_num;
	int cell_ID = (x - 0 ) / cell_length;

	size = cuda_get_cell_size(cell, cell_start, cell_ID, n, cell_num);
	
	int L = 0;
	count = size;
	cell_list_size = 1;
	int* list = (int*)malloc(sizeof(int)*cell_list_size); 
	list[0] = cell_ID;

	while ( count < k )
	{
		//search at level L.
		L++;

		//before allocating more space, I need to know how many more entries are needed, 1 or 2. Depending on whether
		// the neighbor of x is in legal range.
		int more = 0;
		double x1 = x - L*cell_length;
		double x2 = x + L*cell_length;
		if ( x1 > 0 ) more++; 
		if ( x2 < L0 ) more++;

		//enlarge list, keep original data.
		int* tmp_list = (int*)malloc(sizeof(int)*cell_list_size);
		memcpy(tmp_list, list, sizeof(int)*cell_list_size);
		free( list );
		list = (int*)malloc( sizeof(int)*(cell_list_size + more) );
		memcpy( list, tmp_list, sizeof(int)*(cell_list_size) );
		free( tmp_list );

		if ( x1 > 0 ) 
		{
			cell_list_size++;
			cell_ID = (x1 - 0 ) / cell_length;
			size = cuda_get_cell_size(cell, cell_start, cell_ID, n, cell_num);
			count += size;
			list[cell_list_size - 1] = cell_ID;
		}

		if (x2 < L0)
		{
			cell_list_size++;
			cell_ID = (x2 - 0 ) / cell_length;
			size = cuda_get_cell_size(cell, cell_start, cell_ID, n, cell_num);
			count += size;
			list[cell_list_size -1] = cell_ID;
		}
	}
	free ( *cell_list );
	*cell_list = list;
}


__global__ void cuda_k_neighbor_search(double* dx, double* dvx, int* cell, int* ID, int* cell_start,
	int n, int k, int cell_num,
	int* test)
{
	int tid = threadIdx.x + blockIdx.x * blockDim.x;
	if ( tid < n)
	{	
		//cuda_get_cell_size(cell, cell_start, cell[tid], n, cell_num, cell_size);
		if (tid == 0)
		{
			int* cell_list = (int*)malloc(sizeof(int)*1);
			int cell_list_size = 0;
			int count = 0;
			cuda_get_cell_list(cell, cell_start, dx[tid], n, k, cell_num, &cell_list, cell_list_size, count);

			int* neighbor_list = (int*)malloc(sizeof(int)*count);
			int prefix = 0;
			int cell_size = 0;
			for (int i = 0; i < cell_list_size; i++)
			{
				int start = cell_start[cell_list[i]];
				cell_size = cuda_get_cell_size(cell, cell_start, cell_list[i], n, cell_num);
				for (int j = 0; j < cell_size; j++)
					neighbor_list[prefix + j] = ID[start + j];
				prefix += cell_size;
			}


			//neighbor_list[i] is the i-th individual.
			//Warning: there is a transition from tid to ID[tid], need to be carefully examined.
			double* square_dst = (double*)malloc(sizeof(double)*count);
			for (int i = 0; i < count; i++)
				square_dst[i] = (dx[tid] - dx[neighbor_list[i]])*(dx[tid] - dx[neighbor_list[i]]);

			int* k_neighbor_list = (int*)malloc(sizeof(int)*k);
			cuda_findKthMin2(square_dst, count, k, k_neighbor_list);

			//Warning: be careful not to use k_neighbor_list as neighbor id, use neighbor_list[k_neighbor_list] to
			//retrieve k-neighbor id.
			for (int i = 0; i < k; i++)
				test[i] = neighbor_list[k_neighbor_list[i]];


			//Do the update of speed.
			double avg_dvx = 0;
			for (int i = 0; i < k; i++)
				avg_dvx += dvx[neighbor_list[k_neighbor_list[i]]];
			avg_dvx /= k;
			
		}
	}
}
		

void initialize(thrust::device_vector< thrust::pair<int,int> >& v)
{
	thrust::default_random_engine rng(123456);
	thrust::uniform_int_distribution<int> dist(0,9);
	for(size_t i = 0; i < v.size(); i++)
	{
		int a = dist(rng);
		int b = dist(rng);
		v[i] = thrust::make_pair(a,b);
	}
}

void thrust_simple_test()
{
	int n = 32;
	int threadNum = 128;
	int blockNum = (n + threadNum - 1) / threadNum;
	const int cellNum = 5;
	thrust::device_vector<double> dev_x(n);
	thrust::device_vector<double> dev_vx(n);
	srand(1);
	for (int i = 0; i < n; i++)
	{
		dev_x[i] = (double)rand() / RAND_MAX * L0;
		dev_vx[i] = (double)rand() / RAND_MAX * V0;
	}

	double* ptr_x = thrust::raw_pointer_cast(&dev_x[0]);
	double* ptr_vx = thrust::raw_pointer_cast(&dev_vx[0]);

	thrust::device_vector<int> dev_cell(n);
	int* ptr_cell = thrust::raw_pointer_cast(&dev_cell[0]);

	cuda_hash_mapping<<<blockNum, threadNum>>>(ptr_x, ptr_cell, n, 0, L0, cellNum);

	thrust::device_vector<int> dev_ID(n);
	int* ptr_ID = thrust::raw_pointer_cast(&dev_ID[0]);

	cuda_sort_cell(dev_cell, dev_ID);

	thrust::device_vector<int> dev_cellStart(cellNum + 1);
	int* ptr_cellStart = thrust::raw_pointer_cast(&dev_cellStart[0]);

	cuda_count_cell<<<blockNum, threadNum>>>(ptr_cell, ptr_cellStart, n,cellNum);

	int k0 = 15;
	thrust::device_vector<int> dev_test(n);
	int* ptr_test = thrust::raw_pointer_cast(&dev_test[0]);

	thrust::device_vector<double> dev_new_dvx(n);
	double* ptr_new_dvx = thrust::raw_pointer_cast(&dev_new_dvx[0]);
	thrust::device_vector<int> test_size(1);
	int* ptr_test_size = thrust::raw_pointer_cast(&test_size[0]);
	cuda_k_neighbor_search<<<blockNum, threadNum>>>(ptr_x, ptr_vx, ptr_cell, ptr_ID, ptr_cellStart, 
								n, k0, cellNum, ptr_test);

	
	print(dev_cell);
	print(dev_ID);
	print(dev_cellStart);
	print(dev_test);
	//for (int i = 0; i < n; i++)
	//{
	//	for (int j = 0; j < k0; j ++)
	//		cout<<dev_test[i*k0 + j]<<" ";
	//	cout<<endl;
	//}
}

int main()
{
	thrust_simple_test();
	//std::cout << "\nsorting pairs\n";
	//{
	//	thrust::device_vector< thrust::pair<int,int> > keys(N);
	//	initialize(keys);
	//	print(keys);
	//	thrust::sort(keys.begin(), keys.end());
	//	print(keys);
	//}
	//cudaEvent_t start,finish;
	//cudaEventCreate( &start );
	//cudaEventCreate( &finish );
	//cudaEventRecord( start, 0);
	//cudaEventRecord( finish, 0);
	//cudaEventSynchronize( finish );
	//float elapsedTime;
	//cudaEventElapsedTime(&elapsedTime, start, finish);
	//printf("%3.1f ms", elapsedTime/10);
	////cuda_copyandshow(hx, dvx, N);
	//cudaEventDestroy( start );
	//cudaEventDestroy( finish );;
	system("pause");
	return 0;
}